import torch
import torch.nn as nn
import torch.nn.functional as F


class Encoder(nn.Module):
    def __init__(self, hyper_params):
        super(Encoder, self).__init__()
        self.linear1 = nn.Linear(
            hyper_params['rnn_size'] * 2 + 32, hyper_params['hidden_size']
        )
        self.linear2 = nn.Linear(
            3*hyper_params['rnn_size'], hyper_params['hidden_size']
        )
        nn.init.xavier_normal_(self.linear1.weight)
        self.activation = nn.Tanh()

        self.no_cnn = nn.Linear(
            hyper_params['rnn_size'], hyper_params['hidden_size']
        )

        nn.init.xavier_normal_(self.no_cnn.weight)
        # 1 x 1 x seq_len x rnn_size ==> 1 x n x seq_len x rnn_size
        self.Horizontal_cnn = nn.Conv2d(
            in_channels=1,
            out_channels=32,  # n
            kernel_size=(1, hyper_params),
            stride=(1, 1)
        )

        self.relu = nn.ReLU()

        # 1 x 1 x seq_len x rnn_size ==> 1 x 1 x seq_len x rnn_size
        self.Vertical_cnn = nn.Conv2d(
            in_channels=1,  # k
            out_channels=1,
            kernel_size=(4, 1),
            stride=(1, 1),
        )
        self.hyper_params = hyper_params
        self.dropout = nn.Dropout(0.2)

    def forward(self, x):

        # 纵向
        s = torch.zeros(
            [x.shape[0], 1, 3, self.hyper_params['rnn_size']], device=x.device)
        q = torch.cat((s, x.unsqueeze(1)), 2)
        q = self.Vertical_cnn(q)
        q = q * x.unsqueeze(1)  # [batch_size x 1 x seq_len x rnn_size]
        q = q.squeeze(1)

        # 横向
        # [batch_size x n x seq_len x rnn_size]
        p = self.Horizontal_cnn(x.unsqueeze(1))
        p = self.relu(p)
        p = torch.sum(p, dim=3).transpose(1, 2)  # [batch_size x seq_len x n]

        # [batch_size x seq_len x (2rnn_size + 32)]
        m = torch.cat((p, x, q), dim=2)
        x = self.linear1(m)
        x = self.activation(x)
        x = self.dropout(x)

        return x


class Decoder(nn.Module):
    def __init__(self, hyper_params):
        super(Decoder, self).__init__()
        self.linear1 = nn.Linear(
            hyper_params['latent_size'], hyper_params['item_embed_size'])
        self.linear2 = nn.Linear(
            hyper_params['item_embed_size'], hyper_params['total_items'] + 1)
        nn.init.xavier_normal_(self.linear1.weight)
        nn.init.xavier_normal_(self.linear2.weight)
        self.activation = nn.Tanh()

    def forward(self, x):
        x = self.linear1(x)
        out_embed = x

        x = self.activation(x)
        x = self.linear2(x)
        return x, out_embed


class Caser(nn.Module):
    def __init__(self, hyper_params):
        super(Model, self).__init__()
        self.hyper_params = hyper_params

        self.encoder = Encoder(hyper_params)
        self.decoder = Decoder(hyper_params)

        self.item_embed = nn.Embedding(
            hyper_params['total_items'] + 1, hyper_params['item_embed_size'])

        self.gru = nn.GRU(
            hyper_params['item_embed_size'], hyper_params['rnn_size'],
            batch_first=True, num_layers=1
        )

        self.linear_o = nn.Linear(
            hyper_params['hidden_size'], hyper_params['latent_size'])
        self.linear1 = nn.Linear(
            hyper_params['hidden_size'], 2 * hyper_params['latent_size'])
        nn.init.xavier_normal_(self.linear1.weight)

        self.tanh = nn.Tanh()
        self.embed_dropout = nn.Dropout(0.2)
        self.encoded_dropout = nn.Dropout(0.2)

        self.add_eps1 = AddEps(hyper_params['rnn_size'])
        self.add_eps2 = AddEps(hyper_params['hidden_size'])

    def sample_latent(self, z_inferred):
        return torch.randn_like(z_inferred)

    def forward(self, x):
        x = self.item_embed(x)
        x_real = x
        # x = self.add_eps1(x)
        rnn_out, _ = self.gru(self.embed_dropout(x))
        #z_inferred = self.encoder(self.add_eps1(rnn_out))
        z_inferred = self.encoder(rnn_out)
        #z_inferred = self.linear_o(self.add_eps2(z_inferred))
        z_inferred = self.linear_o(z_inferred)
        z_inferred = self.encoded_dropout(z_inferred)
        # [batch_size x seq_len x total_items]
        dec_out, out_embed = self.decoder(z_inferred)

        return dec_out, x_real, z_inferred, out_embed
